import torch

# Load model directly
from transformers import (BertTokenizer,
                          AutoModelForMaskedLM,
                          AutoModel,
                          AutoTokenizer,
                          AutoModelForQuestionAnswering,
                          AutoModelForSequenceClassification,
                          AutoModelForTokenClassification)

from transformers.modeling_outputs import MaskedLMOutput, BaseModelOutputWithPoolingAndCrossAttentions
import os

# 设置离线模式（使用本地缓存）
os.environ['HF_DATASETS_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['HF_ENDPOINT'] = "https://hf-mirror.com"

"""
1. AutoTokenizer
    用途：用于加载预训练模型的 tokenizer。
    功能：将文本转换为模型可以理解的输入格式（如 token IDs、attention masks 等）。
2. AutoModel
    用途：用于加载预训练模型的基础架构（不包含特定任务的输出层）。
    功能：返回模型的隐藏状态（hidden states），适用于自定义任务或特征提取。
3. AutoModelForXXX
    用途：用于加载特定任务的预训练模型（带有任务特定的输出层）。
    常见任务：
        AutoModelForSequenceClassification：用于文本分类。
        AutoModelForTokenClassification：用于序列标注（如命名实体识别）。
        AutoModelForQuestionAnswering：用于问答任务。
        AutoModelForMaskedLM：用于掩码语言模型任务（如 BERT 的填空任务）。
        AutoModelForCausalLM：用于因果语言模型任务（如 GPT 的文本生成）。
4. AutoConfig
    用途：用于加载模型的配置（如隐藏层大小、注意力头数等）。
    功能：可以用于自定义模型或查看模型的配置。
5. Pipeline
    用途：提供简单易用的高级 API，用于常见任务（如文本分类、命名实体识别、文本生成等）。
    功能：封装了 tokenizer 和模型，直接处理原始文本输入。
"""
pretrain_model_dir = 'E:\\ai\\huggingface-models\\bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(pretrain_model_dir,
                                          local_files_only=True, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(pretrain_model_dir,
                                             local_files_only=True, trust_remote_code=True)

# 老的API，已经不适用了
# source = 'huggingface/pytorch-transformers'
# part = 'tokenizer'
# model_name = 'bert-base-chinese'
# tokenizer = torch.hub.load(source, part, model_name)

print(f'tokenizer:{tokenizer}')
# print(f'model:{model}')

input_text = '人生该如何起头'
index_tokens = tokenizer.encode(input_text)

# index_tokens:[101, 782, 4495, 6421, 1963, 862, 6629, 1928, 102]
# 101：句子的开头，102：句子的结尾，这两个是固定值
# 如果传给input_text的多条句子，index_tokens的值可能为：
# index_tokens:[101, 782, 4495, 6421, 102, 1963, 862, 6629, 1928, 102]
# 中间的102为句子的分隔符
print(f'index_tokens:{index_tokens}')
tokens_tensor = torch.tensor([index_tokens])

#
token_output = tokenizer('hello world!')
print(f'token_output: {token_output}')
print(f'token_output.input_ids.shape: {token_output.input_ids}')

with torch.no_grad():
    encoded_layers: MaskedLMOutput = model(tokens_tensor)

print(f'encoded_layers: {encoded_layers.logits.shape}')
